#!/usr/bin/env python
import Crypto.PublicKey.RSA
import Crypto.Util.number
import pyasn1.type.univ
import pyasn1.codec.der.encoder
import pyasn1.codec.der.decoder

class KeyParser(object):
    oid = None  # This should be set by subclasses.
    _algorithms = []

    def __init__(self):
        self._params = None

    @classmethod
    def register(cls):
        cls._algorithms.append((cls.oid, cls))

    @classmethod
    def get_by_oid(cls, oid):
        for tmp_oid, tmp_parser in cls._algorithms:
            if tmp_oid == oid:
                return tmp_parser
        else:
            raise Exception('No parser found for algo OID', oid)

    def set_algo_params(self, params):
        self._params = params

class RsaKeyParser(KeyParser):
    # pkcs-1    OBJECT IDENTIFIER ::= {
    #     iso(1) member-body(2) us(840) rsadsi(113549) pkcs(1) 1
    # }
    # rsaEncryption OBJECT IDENTIFIER ::= { pkcs-1 1 }
    #
    # RSAPublicKey ::= SEQUENCE {
    #     modulus           INTEGER,  -- n
    #     publicExponent    INTEGER   -- e
    # }
    # RSAPrivateKey ::= SEQUENCE {
    #     version           Version,
    #     modulus           INTEGER, -- n
    #     publicExponent    INTEGER, -- e
    #     privateExponent   INTEGER, -- d
    #     prime1            INTEGER, -- p
    #     prime2            INTEGER, -- q
    #     exponent1         INTEGER, -- d mod (p-1)
    #     exponent2         INTEGER, -- d mod (q-1)
    #     coefficient       INTEGER, -- (inverse of q) mod p
    #     otherPrimeInfos   OtherPrimeInfos OPTIONAL
    # }
    # Version ::= INTEGER { two-prime(0), multi(1) }
    #     (CONSTRAINED BY {-- version must be multi if otherPrimeInfos present --})
    #
    # The parameters field associated with this OID in a value of
    # type AlgorithmIdentifier shall have a value of type NULL.

    oid = '1.2.840.113549.1.1.1'

    def parse_key(self, data):
        # convert data to bytes if necessary
        if isinstance(data, pyasn1.type.univ.BitString):
            bits = list(data) + [0] * ((8 - len(data) % 8) % 8)
            bytes = []
            for i in range(0, len(bits), 8):
                bytes.append(chr(int(''.join(str(b) for b in bits[i:i + 8]), 2)))
            data = ''.join(bytes)
        outer, leftover = pyasn1.codec.der.decoder.decode(data)
        assert not leftover
        assert isinstance(outer, pyasn1.type.univ.Sequence)
        if len(outer) == 2:
            # public key (n, e)
            temp = [long(x) for x in outer]
            # construct_py generally returns correct keys so we use it directly
            return Crypto.PublicKey.RSA.construct_py(temp)
        else:
            # private key (n, e, d, p, q, u)
            assert len(outer) >= 9
            temp = [long(x) for x in outer[1:6] + [outer[8]]]
            # construct_py generally returns correct keys so we use it directly
            return Crypto.PublicKey.RSA.construct_py(temp)

RsaKeyParser.register()

def key_from_file(filename):
    return key_from_string(file(filename).read())

def key_from_string(data):
    outer, leftover = pyasn1.codec.der.decoder.decode(data)
    assert not leftover

    assert isinstance(outer, pyasn1.type.univ.Sequence)
    if len(outer) == 2:
        #SubjectPublicKeyInfo format
        algo_identifier, key_data = outer
        assert isinstance(algo_identifier, pyasn1.type.univ.Sequence)
        assert isinstance(key_data, pyasn1.type.univ.BitString)

        # ALGORITHM-IDENTIFIER ::= CLASS {
        #   &id    OBJECT IDENTIFIER UNIQUE,
        #   &Type OPTIONAL
        # }
        # AlgorithmIdentifier { ALGORITHM-IDENTIFIER:InfoObjectSet } ::= SEQUENCE {
        #   algorithm
        #     ALGORITHM-IDENTIFIER.&id({InfoObjectSet}),
        #   parameters
        #     ALGORITHM-IDENTIFIER.&Type({InfoObjectSet}{@.algorithm}) OPTIONAL
        # }
        assert len(algo_identifier) == 2
        algo_oid, algo_params = algo_identifier

        assert isinstance(algo_oid, pyasn1.type.univ.ObjectIdentifier)
        algo_oid_s = '.'.join(str(x) for x in algo_oid)
        parser_cls = KeyParser.get_by_oid(algo_oid_s)
        parser = parser_cls()
        parser.set_algo_params(algo_params)
    else:
        #RSAPrivateKey format
        key_data = data
        parser = RsaKeyParser()

    return parser.parse_key(key_data)

if __name__ == '__main__':
    # openssl genrsa -out foo.key
    # openssl rsa -in foo.key -pubout -outform der -out foo.pub

    key = key_from_file('private.der')
    print key
    #print key.n
    print Crypto.Util.number.long_to_bytes(key.n).encode('hex')
    print key.e
